import logging
import warnings
import os
import base64
import json
import traceback
import tempfile
from configparser import ConfigParser
from concurrent.futures import ThreadPoolExecutor, as_completed
from PIL import Image, UnidentifiedImageError
from pdf2image import convert_from_path
import docx

from maas.constants.http_method_name import HttpMethodName
from maas.http.api_client import ApiClient
from maas.model.api_request import ApiRequest

# --- 日志和警告配置 ---
warnings.filterwarnings("ignore", category=UserWarning)
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S",
                    level=logging.INFO)
logger = logging.getLogger(__name__)

# --- 全局配置变量 ---
OCR_API_URL, OCR_API_KEY, OCR_API_SECRET, OCR_MODEL_NAME = None, None, None, None

# --- 提示词 (Prompts) ---
OCR_PROMPT = """请精确识别图片中的所有文字内容，并严格按照原始的段落格式和顺序返回。"""
TYPE_PROMPT = """请严格判断图片中的主要文字类型是手写体还是印刷体。你的回答只能是 'handwritten' 或 'non-written' 这两个词中的一个，不要包含任何其他解释或标点符号。"""


# --- 配置初始化函数 ---
def init_ocr(config: ConfigParser):
    """初始化OCR服务配置。"""
    global OCR_API_URL, OCR_API_KEY, OCR_API_SECRET, OCR_MODEL_NAME
    try:
        section = 'ocr'
        OCR_API_URL = config.get(section, 'api_url')
        OCR_API_KEY = config.get(section, 'api_key')
        OCR_API_SECRET = config.get(section, 'api_secret')
        OCR_MODEL_NAME = config.get(section, 'model_name')
        logger.info(f"OCR服务配置加载成功, 模型: {OCR_MODEL_NAME}")
    except Exception as e:
        logger.error(f"加载OCR服务配置失败: {e}")
        OCR_API_URL, OCR_API_KEY, OCR_API_SECRET, OCR_MODEL_NAME = None, None, None, None


# --- 图像处理与API调用核心 ---
def _compress_image(input_path: str, output_path: str, max_size_kb: int = 15360):
    """压缩图片以防止API请求体过大。"""
    with Image.open(input_path) as img:
        if img.mode in ("RGBA", "P"):
            img = img.convert("RGB")
        img.thumbnail((2048, 2048), Image.Resampling.LANCZOS)
        quality = 90
        img.save(output_path, "JPEG", quality=quality, optimize=True)
        if os.path.getsize(output_path) / 1024 < max_size_kb:
            logger.info(f"图片 '{os.path.basename(input_path)}' 已优化至 {os.path.getsize(output_path) / 1024:.2f} KB")
            return
        quality = 80
        while quality > 20:
            img.save(output_path, "JPEG", quality=quality, optimize=True)
            if os.path.getsize(output_path) / 1024 < max_size_kb:
                logger.info(
                    f"图片 '{os.path.basename(input_path)}' 已压缩至 {os.path.getsize(output_path) / 1024:.2f} KB")
                return
            quality -= 10
    logger.warning(f"无法将图片压缩到{max_size_kb}KB以下")


def image_to_base64(image_path: str) -> str:
    """将本地图片文件转换为Base64编码。"""
    with open(image_path, "rb") as image_file:
        return f"data:image/jpeg;base64,{base64.b64encode(image_file.read()).decode('utf-8')}"


def call_maas_vlm_api(image_path: str, prompt: str) -> str:
    if not all([OCR_API_URL, OCR_API_KEY, OCR_API_SECRET, OCR_MODEL_NAME]):
        raise RuntimeError("OCR服务配置不完整或未初始化。")

    base64_image = image_to_base64(image_path)
    api_request = ApiRequest(HttpMethodName.POST, OCR_API_URL)
    api_request.add_headers("Content-Type", "application/json")
    api_request.set_credential(OCR_API_KEY, OCR_API_SECRET)
    request_body = {
        "model": OCR_MODEL_NAME,
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": base64_image}},
                {"type": "text", "text": prompt}
            ]}
        ], "max_tokens": 2048, "temperature": 0.1, "stream": False
    }
    api_request.set_json_body(json.dumps(request_body))
    response = ApiClient().send_request(api_request)
    if response.ok:
        return response.json().get("choices", [{}])[0].get("message", {}).get("content", "").strip()
    else:
        raise RuntimeError(f"API请求失败: {response.status_code} - {response.text}")


# --- 文档处理流程 ---
def _process_single_image(image_path: str, source_info: str) -> dict:
    """
    对单个图片进行分类和OCR的完整流程。这是将在线程中运行的函数。
    """
    temp_files_to_clean = []
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpeg") as tmp:
            compressed_path = tmp.name
        temp_files_to_clean.append(compressed_path)
        _compress_image(image_path, compressed_path)

        logger.info(f"正在分类图片来源: {source_info}")
        type_response = call_maas_vlm_api(compressed_path, TYPE_PROMPT).lower()
        is_handwritten = 'handwritten' in type_response and 'non' not in type_response

        if not is_handwritten:
            logger.info(f"来源 '{source_info}' 的图片为非手写体，跳过OCR。")
            return {"source": source_info, "handwritten": False, "text": ""}

        logger.info(f"来源 '{source_info}' 的图片为手写体，开始OCR...")
        ocr_result = call_maas_vlm_api(compressed_path, OCR_PROMPT)
        cleaned_text = ocr_result.replace('\n', ' ').strip()
        final_text = cleaned_text if len(cleaned_text) > 10 else "（本页无有效文本内容）"

        return {"source": source_info, "handwritten": True, "text": final_text}
    except Exception as e:
        logger.error(f"处理来源 '{source_info}' 的图片时失败: {e}")
        return {"source": source_info, "handwritten": False, "text": f"处理失败: {str(e)}"}
    finally:
        for path in temp_files_to_clean:
            if os.path.exists(path):
                os.unlink(path)


def process_document(file_path: str):
    """
    处理主函数，实现“前五张预检”逻辑，并增加图片有效性验证。
    """
    if not os.path.exists(file_path):
        return {"error": f"文件不存在: {file_path}"}

    file_extension = os.path.splitext(file_path)[1].lower()
    temp_image_paths = []
    source_infos = []

    try:
        # --- 步骤1: 提取所有图片，并验证其有效性 ---
        if file_extension in ['.jpg', '.jpeg', '.png', '.bmp', '.gif']:
            temp_image_paths.append(file_path)
            source_infos.append(f"Image file: {os.path.basename(file_path)}")
        else:
            image_extractor = None
            if file_extension == '.pdf':
                logger.info(f"开始从PDF文件提取图片: {file_path}")
                image_extractor = enumerate(convert_from_path(file_path))
            elif file_extension == '.docx':
                logger.info(f"开始从Word文件提取图片: {file_path}")
                doc = docx.Document(file_path)
                image_extractor = enumerate(
                    rel.target_part.blob for rel in doc.part.rels.values() if "image" in rel.target_ref)

            if image_extractor:
                for i, img_data in image_extractor:
                    source_info = f"PDF Page {i + 1}" if file_extension == '.pdf' else f"DOCX Image {i + 1}"
                    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
                        page_path = tmp.name
                    temp_image_paths.append(page_path)

                    try:
                        if isinstance(img_data, Image.Image):  # 来自pdf2image
                            img_data.save(page_path, format="PNG")
                        else:  # 来自docx
                            with open(page_path, 'wb') as img_file:
                                img_file.write(img_data)

                        # 验证图片文件是否有效
                        with Image.open(page_path) as img:
                            img.verify()  # 验证图片数据完整性
                        source_infos.append(source_info)

                    except (IOError, UnidentifiedImageError, Exception) as e:
                        logger.warning(f"跳过无效或损坏的图片: {source_info}。错误: {e}")
                        # 无需添加到source_infos，这张图片将被忽略
                        pass
            else:
                return {"error": f"不支持的文件类型: {file_extension}"}

        if not temp_image_paths:
            logger.info("在文档中未找到任何有效图片。")
            return []

        # 将有效的任务打包
        tasks = list(zip(temp_image_paths, source_infos))

        # --- 步骤2: 对前五张有效图片进行批量预检 ---
        pre_check_tasks = tasks[:5]
        logger.info(f"开始对前 {len(pre_check_tasks)} 张有效图片进行批量预检...")

        pre_check_results = []
        with ThreadPoolExecutor(max_workers=5) as executor:
            future_to_task = {executor.submit(_process_single_image, path, info): info for path, info in
                              pre_check_tasks}
            for future in as_completed(future_to_task):
                pre_check_results.append(future.result())

        # --- 步骤3: 检查预检结果 ---
        for result in pre_check_results:
            if not result.get("handwritten"):
                logger.warning(f"预检失败：来源 '{result['source']}' 被判断为非手写体。已终止整个文档的处理。")
                return [
                    {"source": result['source'], "handwritten": False, "text": "（预检发现非手写体图片，已跳过整个文档）"}]

        # --- 步骤4: 如果预检通过，则并发处理剩余图片 ---
        remaining_tasks = tasks[5:]
        if not remaining_tasks:
            logger.info("所有图片均已在预检中处理完毕。")
            pre_check_results.sort(key=lambda x: x['source'])
            return pre_check_results

        logger.info(f"预检通过，将对剩余的 {len(remaining_tasks)} 张图片进行OCR。")
        remaining_results = []
        with ThreadPoolExecutor(max_workers=5) as executor:
            future_to_task = {executor.submit(_process_single_image, path, info): info for path, info in
                              remaining_tasks}
            for future in as_completed(future_to_task):
                remaining_results.append(future.result())

        final_results = pre_check_results + remaining_results
        final_results.sort(key=lambda x: x['source'])
        return final_results

    finally:
        logger.info("正在清理所有提取出的临时文件...")
        for path in temp_image_paths:
            if path != file_path and os.path.exists(path):
                try:
                    os.unlink(path)
                except Exception as e:
                    logger.error(f"删除临时文件失败: {path}, 错误: {e}")


# --- 主程序入口（用于独立测试） ---
if __name__ == "__main__":
    def test_document_processing():
        print("===== 开始测试文档处理功能 =====")
        try:
            config = ConfigParser()
            if not os.path.exists('config.ini'):
                print("错误：未找到 config.ini 文件。")
                return
            config.read('config.ini', encoding='utf-8')
            init_ocr(config)
        except Exception as e:
            print(f"初始化配置失败: {e}")
            return
        if not OCR_API_URL:
            print("OCR配置加载不成功，测试终止。")
            return

        # --- 在这里修改为您要测试的文件路径 ---
        test_file_path = r"C:\Users\19645\OneDrive\Desktop\AI OKMS_平台用户手册_v1.2.docx"

        if not os.path.exists(test_file_path):
            print(f"错误: 测试文件 '{test_file_path}' 不存在。")
            return

        print(f"\n正在处理文件: {test_file_path}")
        # 引入time模块用于计时
        import time
        start_time = time.time()
        final_result = process_document(test_file_path)
        end_time = time.time()
        print(f"处理完成，总耗时: {end_time - start_time:.2f} 秒")

        print("\n===== 测试结果 =====")
        # 现在的 final_result 直接就是列表
        print(json.dumps(final_result, indent=4, ensure_ascii=False))
        print("====================")


    test_document_processing()